import sys
sys.path.append('/data/xingshiyou-slurm/codeLLM/GPT-respond')

from constant import test_data_path

from data_w_r import read_data, write_data

import matplotlib.pyplot as plt

# 设置 Matplotlib 后端为 Agg，以便在没有显示界面的环境中运行
plt.switch_backend('Agg')

def get_show(lengths, filename):
    # 绘制直方图
    plt.figure(figsize=(10, 6))
    plt.hist(lengths, bins=range(0, max(lengths) + 1000, 1000), color='skyblue', edgecolor='black')

    # 添加标题和标签
    plt.title("Length Distribution Histogram")
    plt.xlabel("Length")
    plt.ylabel("Frequency")

    # 设置横坐标刻度为1000的间隔
    plt.xticks(range(0, max(lengths) + 1000, 1000))

    # 旋转刻度标签，避免重叠
    plt.xticks(rotation=45)

    # 保存图像为文件，指定路径和文件格式
    plt.savefig(filename, format="png")

    # 显示图形
    # plt.show()

def calculate_average(numbers):
    if not numbers:  # 检查列表是否为空
        return 0
    return sum(numbers) / len(numbers)

if __name__ == '__main__':
    test_data = read_data('/data/xingshiyou-slurm/codeLLM/data/test_B/Q_B_without_answer.jsonl')
    prefix_lens = [len(i['prefix']) for i in test_data]
    fim_suffix_lens = [len(i['fim_suffix']) for i in test_data]

    print('前缀最大长度:', max(prefix_lens))
    print('后缀最大长度:', max(fim_suffix_lens))

    print('前缀平均长度:', calculate_average(prefix_lens))
    print('后缀平均长度:', calculate_average(fim_suffix_lens))

    total = [pre + suf for pre, suf in zip(prefix_lens, fim_suffix_lens)]

    get_show(prefix_lens, '../count/prefix_B_len.png')
    get_show(fim_suffix_lens, '../count/suffix_B_len.png')

    get_show(total, '../count/total_B_len.png')
